#!/usr/bin/env bash
# Copyright (c) 2024 maminjie <canpool@163.com>
# SPDX-License-Identifier: MulanPSL-2.0

ROOT_DIR=$(dirname $(readlink -f "$0"))
CUR_DIR=$(pwd)
WORK_DIR=$CUR_DIR/iso_make_workspace

BUILD_DIR="$WORK_DIR/build"
STD_DIR="$BUILD_DIR/standard"
EVE_DIR="$BUILD_DIR/everything"
EVE_SRC_DIR="$BUILD_DIR/everything_src"
EVE_DBG_DIR="$BUILD_DIR/everything_debug"
ISO_DIR="$BUILD_DIR/iso"

YUM_DIR="$WORK_DIR/yum"
YUM_CONF="$YUM_DIR/yum.conf"
YUM_INSTALLROOT="$YUM_DIR/install"

OUT_DIR=$WORK_DIR

ARCH="$(uname -m)"
PRODUCT="MyOS"
VERSION="1.0"
RELEASE=""
NAME="standard"
REPOS=""
ISO_TYPE="standard"
KS_FILE=""
BASEOS_DIR="$ROOT_DIR/os/openEuler"
LORAX_SHAREDIR="$ROOT_DIR/80-openeuler"
CONTINUE_FLAG=0

TYPES="standard everything everything_debug everything_src"

umask 0022


usage() {
printf "make iso image based on yum repository

usage:
    bash iso_make.sh [-h] [-t TYPE] [-p PRODUCT] [-v VERSION] [-r RELEASE] [-s REPOSITORY]
                     [-k KSFILE] [-b BASEOS] [-o OUTDIR] [-n NAME] [-c]

optional arguments:
    -t TYPE        ISO Type, default is standard, includes:
                       ${TYPES}
    -p PRODUCT     product name, such as: openEuler, default is MyOS
    -v VERSION     version identifier, default is 1.0
    -r RELEASE     release information, default is null
    -s REPOSITORY  source dnf repository address link (can be specified multiple times)
    -k KSFILE      kickstart filepath
    -b BASEOS      base os directory that contains normal.xml and rpmlist.xml, default is os/openEuler
    -o OUTDIR      iso output directory, default is work_dir
    -n NAME        field name in iso file, default is standard (currently only applies to standard iso)
    -c             keep the cache and continue the make, not recreate install image
    -h             show the help message and exit
\n"
}


log_error() {
    echo "ERROR: $1"
}

log_info() {
    echo "INFO: $1"
}

check_type() {
    for typename in $TYPES; do
        [[ "${typename}" = "${ISO_TYPE}" ]] && return 0
    done

    return 1
}

parse_cmdline() {
    if [ $# -eq 0 ]; then
        usage; exit 0
    fi

    # parse cmd line arguments
    while getopts ":p:v:r:s:t:k:b:o:n:ch" opt; do
        case "$opt" in
            p)
                PRODUCT="$OPTARG"
            ;;
            v)
                VERSION="$OPTARG"
            ;;
            r)
                RELEASE="$OPTARG"
            ;;
            s)
                if [ -z "$REPOS" ]; then
                    REPOS="$OPTARG"
                else
                    REPOS="$REPOS $OPTARG"
                fi
            ;;
            t)
                ISO_TYPE="$OPTARG"
            ;;
            k)
                KS_FILE="$(realpath $OPTARG)"
            ;;
            b)
                if [ -d "$OPTARG" ]; then
                    BASEOS_DIR="$(realpath $OPTARG)"
                else
                    log_error "no such directory: $OPTARG"
                fi
            ;;
            o)
                OUT_DIR="$(realpath $OPTARG)"
            ;;
            n)
                NAME="$OPTARG"
            ;;
            c)
                CONTINUE_FLAG=1
            ;;
            h)
                usage; exit 0
            ;;
            ?)
                log_error "please check the params."; usage; return 1
            ;;
        esac
    done

    if [ -z "$REPOS" ]; then
        log_error "invalid REPOSITORY arguments"; return 1
    fi

    check_type
    if [ $? -ne 0 ]; then
        log_error "unsupported iso type: ${ISO_TYPE}"; return 1
    fi

    return 0
}

check_user() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root!"
        exit 1
    fi
}

check_arch() {
    if [[ "${ARCH}" != "aarch64" && "${ARCH}" != "x86_64" && "${ARCH}" != "loongarch64" ]]; then
        log_error "unsupported architecture: ${ARCH}"
        exit 1
    fi
}

check_requires() {
    local ret=0

    # createrepo_c package contains:
    #   createrepo
    # dnf-plugins-core package contains:
    #   repoclosure yumdownloader
    # genisoimage package contains:
    #   mkisofs
    # isomd5sum package contains
    #   implantisomd5
    # eg:
    #   dnf install createrepo_c dnf-plugins-core genisoimage isomd5sum -y
    for c in lorax yumdownloader mkisofs createrepo implantisomd5; do
        command -v $c > /dev/null
        if [ $? -ne 0 ]; then
            log_error "command $c not found"
            ret=1
        fi
    done

    return $ret
}

# make_iso iso_label iso_file iso_dir
# params:
#   iso_label - equal to LABEL in EFI/BOOT/grub.cfg
#   iso_file - new iso file name
#   iso_dir - direcotry of iso files
make_iso() {
    if [ $# -lt 3 ]; then
        echo "invalid param"; return 1
    fi

    local arch=$(uname -m)
    local iso_label=$1
    local iso_file=$2
    local iso_dir=$(realpath $3)

    if [ ! -d "$iso_dir" ]; then
        echo "$iso_dir: no such directory"; return 1
    fi

    if [[ "$arch" = "x86_64" ]]; then
        mkisofs -R -J -T -r -l -d -joliet-long -allow-multidot -allow-leading-dots -no-bak -V "${iso_label}" -o "${iso_file}" -b isolinux/isolinux.bin \
            -c isolinux/boot.cat -no-emul-boot -boot-load-size 4 -boot-info-table  -eltorito-alt-boot -e images/efiboot.img -no-emul-boot "${iso_dir}"
        [[ $? -ne 0 ]] && return 1
        isohybrid -u "$iso_file"
    elif [[ "$arch" = "aarch64" || "$arch" = "loongarch64" ]]; then
        mkisofs -R -J -T -r -l -d -joliet-long -allow-multidot -allow-leading-dots -no-bak -V "${iso_label}" -o "${iso_file}" -e images/efiboot.img -no-emul-boot "${iso_dir}"
        [[ $? -ne 0 ]] && return 1
    else
        echo "unsupported architecture: ${arch}"; return 1
    fi
    implantisomd5 "$iso_file"
    return 0
}

# create_yum_repo YUM_DIR YUM_CONF REPOS
create_yum_repo() {
    local yum_dir="$1"
    local yum_config_file="$2"
    local repos="$3"
    local yum_repo_dir="${yum_dir}/yum.repos.d"
    local yum_repo_file="${yum_repo_dir}/os.repo"

    mkdir -p "$yum_repo_dir"

    cat > "${yum_config_file}" <<EOF
[main]
gpgcheck=1
installonly_limit=3
clean_requirements_on_remove=True
best=True
skip_if_unavailable=False
installroot=${YUM_INSTALLROOT}
reposdir=${yum_repo_dir}
EOF

    local i=1
    for repo in $repos; do
        cat >> "${yum_repo_file}" <<EOF
[os_repo_$i]
name=$repo
baseurl=$repo
enabled=1
gpgcheck=0

EOF
        ((i++))
    done
}

create_install_img() {
    local repo_args=""
    for rep in $REPOS; do
        repo_args="${repo_args} -s ${rep}"
    done

    local vertmp="${VERSION}"
    if [ -n "${RELEASE}" ]; then
        vertmp="${VERSION}-${RELEASE}"
    fi
    local lorax_log="lorax.log"

    lorax --isfinal -p "${PRODUCT}" -v "${vertmp}" -r "${RELEASE}" -t "Server" --sharedir $LORAX_SHAREDIR \
        --rootfs-size=4 --buildarch="$ARCH" $repo_args --nomacboot --noupgrade --logfile $lorax_log "${ISO_DIR}"

    if [ $? -ne 0 ] ; then
        cat $lorax_log
        return 1
    fi
    return 0
}

init_dir() {
    [[ ! -d "$1" ]] && mkdir -p "$1"
}

clean_dir() {
    [[ -d "$1" ]] && rm -rf "$1"
}

init_environment() {
    cd $CUR_DIR

    if [ $CONTINUE_FLAG -eq 0 ]; then
        clean_dir $WORK_DIR
    fi
    init_dir $WORK_DIR

    init_dir $BUILD_DIR
    init_dir $YUM_DIR
    init_dir $YUM_INSTALLROOT

    init_dir $OUT_DIR

    cd $WORK_DIR
}

deinit_environment() {
    cd $CUR_DIR
}

init_variable() {
    CONFIG_RPM_LIST="$BASEOS_DIR/rpmlist.xml"
    CONFIG_PACKAGES_LIST_FILE="$BASEOS_DIR/$ARCH/normal.xml"

    if [ ! -f "$CONFIG_RPM_LIST" ]; then
        log_error "$CONFIG_RPM_LIST: no such file"
        return 1
    fi

    if [ ! -f "$CONFIG_PACKAGES_LIST_FILE" ]; then
        log_error "$CONFIG_PACKAGES_LIST_FILE: no such file"
        return 1
    fi

    if [ ! -d "$LORAX_SHAREDIR" ]; then
        log_error "$LORAX_SHAREDIR: no such directory"
        return 1
    fi

    CONFIG="$CONFIG_PACKAGES_LIST_FILE"

    if [ -n "${RELEASE}" ]; then
        RELEASE_NAME="${PRODUCT}-${VERSION}-${RELEASE}-${ARCH}"
        if [ X$NAME = X ]; then
            STD_ISO_NAME="${PRODUCT}-${VERSION}-${RELEASE}-${ARCH}-dvd.iso"
        else
            STD_ISO_NAME="${PRODUCT}-${VERSION}-${RELEASE}-${NAME}-${ARCH}-dvd.iso"
        fi
        EVE_ISO_NAME="${PRODUCT}-${VERSION}-${RELEASE}-everything-${ARCH}-dvd.iso"
        EVE_DBG_ISO_NAME="${PRODUCT}-${VERSION}-${RELEASE}-everything-debug-${ARCH}-dvd.iso"
        EVE_SRC_ISO_NAME="${PRODUCT}-${VERSION}-${RELEASE}-everything-source-dvd.iso"
    else
        RELEASE_NAME="${PRODUCT}-${VERSION}-${ARCH}"
        if [ X$NAME = X ]; then
            STD_ISO_NAME="${PRODUCT}-${VERSION}-${ARCH}-dvd.iso"
        else
            STD_ISO_NAME="${PRODUCT}-${VERSION}-${NAME}-${ARCH}-dvd.iso"
        fi
        EVE_ISO_NAME="${PRODUCT}-${VERSION}-everything-${ARCH}-dvd.iso"
        EVE_DBG_ISO_NAME="${PRODUCT}-${VERSION}-everything-debug-${ARCH}-dvd.iso"
        EVE_SRC_ISO_NAME="${PRODUCT}-${VERSION}-everything-source-dvd.iso"
    fi

    return 0
}

parse_rpmlist_xml() {
    local packagetype=$1
    xmllint --xpath "//packagelist[@type='${packagetype}']/node()" "${CONFIG_RPM_LIST}" > origin_rpmlist_${packagetype}
    cat origin_rpmlist_${packagetype} | grep packagereq | cut -d ">" -f 2 | cut -d "<" -f 1 > parsed_rpmlist_${packagetype}
    return 0
}

download_standard_rpms() {
    cat "${CONFIG}" | grep packagereq | cut -d ">" -f 2 | cut -d "<" -f 1 > _all_rpms.lst
    parse_rpmlist_xml "${ARCH}"
    cat parsed_rpmlist_${ARCH} >> _all_rpms.lst
    parse_rpmlist_xml "common"
    cat parsed_rpmlist_common >> _all_rpms.lst
    sort -r -u _all_rpms.lst -o _all_rpms.lst

    ret=0
    rm -rf not_find __*
    yum list -c "${YUM_CONF}" available | awk '{print $1}' > __list.arch

    set +e
    local releasever=$(repoquery -c "${YUM_CONF}" -q -i --whatprovides system-release | grep "^Version" | head -1 | awk '{print $NF}')
    for rname in $(cat _all_rpms.lst); do
        if [ -z "${rname}" ]; then
            continue
        fi
        if [ $(echo "$rname" | grep "\*$") ]; then
            echo "$rname" >> __rpm.list
            continue
        fi
        rarch=${rname##*.}
        if [ "X$rarch" == "Xi686" ] && [ "$ARCH" == "aarch64" ]; then
            continue
        fi
        if [ "X$rarch" == "Xx86_64" ] && [ "$ARCH" == "aarch64" ]; then
            rname=${rname%%.*}
            rarch="aarch64"
        fi
        cat __list.arch | grep -w "^$rname" > /dev/null 2>&1
        if [ $? != 0 ]; then
            rname_tmp=`repoquery -c "${YUM_CONF}" --queryformat="%{name}.%{arch}" -q --whatprovides $rname --releasever=${releasever}`
            if [ -z "${rname_tmp}" ]; then
                echo "cannot find $rname in yum repo" >> not_find
                ret=1
                continue
            else
                echo "$rname" >> __rpm.list
                continue
            fi
            rname=${rname_tmp}
        fi
        if [ "X$rarch" == "Xi686" ] || [ "X$rarch" == "Xx86_64" ] || [ "X$rarch" == "Xnoarch" ] || [ "X$rarch" == "Xaarch64" ]; then
            rname="${rname}"
        else
            cat __list.arch | grep -w "^$rname.$ARCH" > /dev/null 2>&1
            if [ $? == 0 ]; then
                rname="${rname}"."${ARCH}"
            else
                cat __list.arch | grep -w "^$rname.noarch" > /dev/null 2>&1
                if [ $? == 0 ]; then
                    rname="${rname}".noarch
                else
                    echo "cannot find $rname in yum repo" >> not_find
                    ret=1
                fi
            fi
        fi
        echo "$rname" >> __rpm.list
    done
    if [ "${ret}" -ne 0 ]; then
        cat not_find|sort|uniq
        exit "${ret}"
    fi

    parse_rpmlist_xml "exclude"
    local exclude_cmd=""
    if [ -s parsed_rpmlist_exclude ]; then
        for rpmname in $(cat parsed_rpmlist_exclude);do
            exclude_cmd="${exclude_cmd} -x ${rpmname}"
        done
    fi
    local yumdownloader_log_startline=$(($(awk 'END{print NR}' /var/log/dnf.log)+1))
    yumdownloader --resolve -c "${YUM_CONF}" --destdir="${STD_DIR}" $(cat __rpm.list | tr '\n' ' ') ${exclude_cmd}
    if [ $? != 0 ] || sed -n ''${yumdownloader_log_startline}',$p' "${YUM_INSTALLROOT}"/var/log/dnf.log | grep -n 'conflicting requests'; then
       echo "Download rpms failed!"
       exit 133
    fi

    parse_rpmlist_xml "conflict"
    set -e
    if [ -s parsed_rpmlist_conflict ]; then
        yumdownloader --resolve -c "${YUM_CONF}" --destdir="${STD_DIR}" $(cat parsed_rpmlist_conflict | tr '\n' ' ') ${exclude_cmd}
    fi
}

download_rpms() {
    if [ $CONTINUE_FLAG -eq 0 ]; then
        clean_dir $YUM_INSTALLROOT
    fi

    set +e
    if [ "${ISO_TYPE}" == "standard" ]; then
        download_standard_rpms
    elif [ "${ISO_TYPE}" == "everything" ]; then
        download_everything_rpms
    elif [ "${ISO_TYPE}" == "everything_src" ]; then
        download_everything_source_rpms
    elif [ "${ISO_TYPE}" == "everything_debug" ]; then
        download_everything_debug_rpms
    fi

    return 0
}

get_rpm_pub_key() {
    # clean old files
    rm -f $ISO_DIR/RPM-GPG-KEY-*
    ls $ISO_DIR/Packages/*-gpg-keys* &> /dev/null
    if [ $? -ne 0 ]; then
        return
    fi
    mkdir -p $ISO_DIR/GPG_tmp
    cp $ISO_DIR/Packages/*-gpg-keys* $ISO_DIR/GPG_tmp
    cd $ISO_DIR/GPG_tmp
    rpm2cpio *-gpg-keys* | cpio -div ./etc/pki/rpm-gpg/RPM-GPG-KEY-*
    cd -
    cp $ISO_DIR/GPG_tmp/etc/pki/rpm-gpg/RPM-GPG-KEY-* $ISO_DIR
    rm -rf $ISO_DIR/GPG_tmp
}

get_everything_rpms() {
    yum list -c "${YUM_CONF}" --available | awk '{print $1}' | grep -E "\.noarch|\.${ARCH}" | grep -vw "debuginfo" | grep -vw "debugsource" > ava_every_lst
    parse_rpmlist_xml "exclude"
    cat parsed_rpmlist_exclude
    if [ -s parsed_rpmlist_exclude ]; then
        for rpmname in $(cat parsed_rpmlist_exclude); do
            sed -i "/^${rpmname}\./d" ava_every_lst
        done
    fi
    if [ -s conflict_list ]; then
        rm -rf conflict_list
    fi
    parse_rpmlist_xml "conflict"
    cat parsed_rpmlist_conflict
    if [ -s parsed_rpmlist_conflict ]; then
        for rpmname in $(cat parsed_rpmlist_conflict); do
            cat ava_every_lst | grep "^${rpmname}\."
            if [ $? -eq 0 ]; then
                sed -i "/^${rpmname}\./d" ava_every_lst
                echo "${rpmname}" >> conflict_list
            fi
        done
    fi
    parse_rpmlist_xml "everything_conflict"
    cat parsed_rpmlist_everything_conflict
    if [ -s parsed_rpmlist_everything_conflict ]; then
        for rpmname in $(cat parsed_rpmlist_everything_conflict); do
            cat ava_every_lst | grep "^${rpmname}\."
            if [ $? -eq 0 ]; then
                sed -i "/^${rpmname}\./d" ava_every_lst
                echo "${rpmname}" >> conflict_list
            fi
        done
    fi
}

download_everything_rpms() {
    mkdir -p ${EVE_DIR}
    get_everything_rpms
    yumdownloader --resolve -c "${YUM_CONF}" --destdir="${EVE_DIR}" $(cat ava_every_lst | tr '\n' ' ')
    if [ $? != 0 ] || [ $(ls ${EVE_DIR} | wc -l) == 0 ]; then
       echo "Download rpms failed!"
       exit 133
    fi
    if [ -s conflict_list ]; then
        yumdownloader --resolve -c "${YUM_CONF}" --destdir="${EVE_DIR}" $(cat conflict_list | tr '\n' ' ')
    fi
}

download_everything_source_rpms() {
    mkdir -p ${EVE_SRC_DIR}
    yum list -c "${YUM_CONF}" --available | awk '{print $1}' | grep ".src" > ava_every_lst
    parse_rpmlist_xml "src_exclude"
    cat parsed_rpmlist_src_exclude
    if [ -s parsed_rpmlist_src_exclude ]; then
        for rpmname in $(cat parsed_rpmlist_src_exclude); do
            sed -i "/^${rpmname}\./d" ava_every_lst
        done
    fi
    yumdownloader -c "${YUM_CONF}" --destdir="${EVE_SRC_DIR}" --source $(cat ava_every_lst | tr '\n' ' ')
    if [ $? != 0 ] || [ $(ls ${EVE_SRC_DIR} | wc -l) == 0 ]; then
       echo "Download rpms failed!"
       exit 133
    fi
}

download_everything_debug_rpms() {
    mkdir -p ${EVE_DBG_DIR}
    yum list -c "${YUM_CONF}" --available | awk '{print $1}' | grep -Ew "debuginfo|debugsource" > ava_debug_lst
    parse_rpmlist_xml "everything_debug_exclude"
    cat parsed_rpmlist_everything_debug_exclude
    if [ -s parsed_rpmlist_everything_debug_exclude ]; then
        for rpmname in $(cat parsed_rpmlist_everything_debug_exclude); do
            sed -i "/^${rpmname}\./d" ava_debug_lst
        done
    fi
    yumdownloader --resolve -c "${YUM_CONF}" --destdir="${EVE_DBG_DIR}" $(cat ava_debug_lst | tr '\n' ' ')
    if [ $? != 0 ] || [ $(ls ${EVE_DBG_DIR} | wc -l) == 0 ]; then
        echo "yumdownloader with --resolve failed, trying to yumdownloader without --resolve"
        yumdownloader -c "${YUM_CONF}" --destdir="${EVE_DBG_DIR}" $(cat ava_debug_lst | tr '\n' ' ')
    fi
}

gen_iso() {
    local iso_name="$STD_ISO_NAME"
    local pkg_dir="$STD_DIR"

    if [ "${ISO_TYPE}" == "everything" ]; then
        iso_name="$EVE_ISO_NAME"
        pkg_dir="$EVE_DIR"
    elif [ "${ISO_TYPE}" == "everything_src" ]; then
        iso_name="$EVE_SRC_ISO_NAME"
        pkg_dir="$EVE_SRC_DIR"
    elif [ "${ISO_TYPE}" == "everything_debug" ]; then
        iso_name="$EVE_DBG_ISO_NAME"
        pkg_dir="$EVE_DBG_DIR"
    fi

    rm -rf "${ISO_DIR}"/Packages
    mv "${pkg_dir}" "${ISO_DIR}"/Packages
    get_rpm_pub_key

    mkdir -p $ISO_DIR/repodata
    rm -rf "${ISO_DIR}"/repodata/*
    cp "$CONFIG" "${ISO_DIR}"/repodata/
    createrepo -g "${ISO_DIR}"/repodata/*.xml "${ISO_DIR}"

    make_iso "${RELEASE_NAME}" "${OUT_DIR}/${iso_name}" "${ISO_DIR}"

    if [ $? -ne 0 ]; then
        log_error "create $ISO_TYPE iso failed"
        return 1
    fi

    ls "${OUT_DIR}/${iso_name}"

    return 0
}

init_config() {
    if [ -f "$KS_FILE" ]; then
        mkdir -p "${ISO_DIR}"/ks
        cp "$KS_FILE" "${ISO_DIR}"/ks/
    else
        clean_dir "${ISO_DIR}"/ks
    fi
    return 0
}

iso_maker() {
    log_info "Init environment..."
    init_environment

    init_variable
    if [ $? -ne 0 ]; then
        return 1
    fi

    if [[ $CONTINUE_FLAG -eq 0 || ! -f "$ISO_DIR/images/install.img" ]]; then
        log_info "Creating install image..."
        create_install_img
        if [ $? -ne 0 ]; then
            log_error "create install image failed"
            return 1
        fi
    fi

    log_info "Creating repos..."
    create_yum_repo "$YUM_DIR" "$YUM_CONF" "$REPOS"
    yum -c "$YUM_CONF" clean all

    log_info "Downloading rpms..."
    download_rpms
    if [ $? -ne 0 ]; then
        log_error "down rpms failed"
        return 1
    fi

    wait

    set +e

    log_info "Initializing config..."
    init_config
    if [ $? -ne 0 ]; then
        log_error "init config failed"
        return 1
    fi

    log_info "Creating iso image..."
    gen_iso
    if [ $? -ne 0 ]; then
        log_error "create iso failed"
        return 1
    fi

    deinit_environment

    return 0
}

trap deinit_environment TERM ABRT EXIT

main() {
    check_user
    check_arch

    parse_cmdline "$@"
    if [ $? -ne 0 ]; then
        log_error "parse params failed, try -h for help"
        return 1
    fi

    check_requires
    if [ $? -ne 0 ]; then
        log_error "check requires failed"
        return 1
    fi

    iso_maker
    if [ $? -ne 0 ]; then
        log_error "make iso failed"
        deinit_environment
        return 1
    fi

    return 0
}

main "$@"
